package com.nanohadoop.yarn;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * 资源调度器，负责协调资源分配
 * ”基本FIFO调度器" （Basic FIFO Scheduler）是一个"先进先出单容器调度算法" （FIFO Single Container Scheduling Algorithm）
 * ## 算法特点
 * - 实现简单，易于理解
 * - 严格按照提交顺序处理应用
 * - 每次调度只为一个应用分配一个容器
 * - 节点选择采用首次适应策略（First Fit）
 * - 不考虑应用优先级或资源使用效率优化
 *
 */
public class FIFOResourceScheduler implements ResourceScheduler {
    private static final Logger log = LoggerFactory.getLogger(FIFOResourceScheduler.class);

    private final Map<String, NodeManager> nodeManagers;
    private final Queue<ApplicationMaster> applicationQueue;
    private final Map<String, ApplicationMaster> applications;
    private int nextContainerId;
    private int nextApplicationId;
    private final ScheduledExecutorService schedulerService;
    private volatile boolean running;

    public FIFOResourceScheduler() {
        this.nodeManagers = new ConcurrentHashMap<>();
        this.applicationQueue = new LinkedList<>();
        this.applications = new ConcurrentHashMap<>();
        this.nextContainerId = 0;
        this.nextApplicationId = 0;
        this.schedulerService = Executors.newSingleThreadScheduledExecutor();
        this.running = false;
        log.debug("FIFOResourceScheduler initialized");
    }

    /**
     * 启动调度器
     */
    @Override
    public void start() {
        if (running) {
            return;
        }
        running = true;
        // 每100毫秒执行一次调度
        schedulerService.scheduleAtFixedRate(this::schedule, 0, 100, TimeUnit.MILLISECONDS);
        log.debug("FIFOResourceScheduler started");
    }

    /**
     * 停止调度器
     */
    @Override
    public void stop() {
        if (!running) {
            return;
        }
        running = false;
        schedulerService.shutdown();
        try {
            if (!schedulerService.awaitTermination(5, TimeUnit.SECONDS)) {
                schedulerService.shutdownNow();
            }
        } catch (InterruptedException e) {
            schedulerService.shutdownNow();
        }
        log.debug("FIFOResourceScheduler stopped");
    }

    /**
     * 注册节点管理器
     */
    @Override
    public void registerNodeManager(NodeManager nodeManager) {
        nodeManagers.put(nodeManager.getNodeId(), nodeManager);
        log.debug("NodeManager registered: " + nodeManager.getNodeId());
    }

    /**
     * 提交应用程序
     */
    @Override
    public String submitApplication(String name, String user, Resource requiredResource, int numContainers) {
        String applicationId = "app_" + nextApplicationId++;
        ApplicationMaster application = new ApplicationMaster(applicationId, name, user, requiredResource, numContainers);
        applications.put(applicationId, application);
        applicationQueue.add(application);
        log.debug("Application submitted: " + applicationId + ", name: " + name + ", user: " + user + ", required resources per container: " + requiredResource.getMemoryMB() + "MB/" + requiredResource.getCpuCores() + "cores, total containers: " + numContainers);
        return applicationId;
    }

    /**
     * 资源调度的核心方法
     */
    @Override
    public void schedule() {
        if (!running) {
            log.debug("Scheduler is not running, skipping schedule");
            return;
        }

        if (applicationQueue.isEmpty()) {
            log.debug("No applications in queue, skipping schedule");
            return;
        }

        if (nodeManagers.isEmpty()) {
            log.debug("No node managers available, skipping schedule");
            return;
        }

        // FIFO策略：按提交顺序处理应用
        ApplicationMaster application = applicationQueue.peek();
        log.debug("Scheduling application: " + application.getApplicationId() + ", current status: " + application.getStatus());

        if (application.getStatus() != ApplicationMaster.ApplicationStatus.SUBMITTED &&
            application.getStatus() != ApplicationMaster.ApplicationStatus.RUNNING) {
            log.debug("Application " + application.getApplicationId() + " has status " + application.getStatus() + ", removing from queue");
            applicationQueue.poll();
            return;
        }

        // 检查是否还需要分配容器
        int allocatedContainers = application.getAllocatedContainers().size();
        int numContainers = application.getNumContainers();
        log.debug("Application " + application.getApplicationId() + ": allocated " + allocatedContainers + "/" + numContainers + " containers");

        if (allocatedContainers >= numContainers) {
            log.debug("All containers allocated for application " + application.getApplicationId() + ", setting status to RUNNING");
            application.setStatus(ApplicationMaster.ApplicationStatus.RUNNING);
            applicationQueue.poll();
            return;
        }

        Resource requiredResource = application.getRequiredResource();
        log.debug("Application " + application.getApplicationId() + " requires " + requiredResource.getMemoryMB() + "MB memory and " + requiredResource.getCpuCores() + " CPU cores per container");

        // 查找有足够资源的节点
        boolean foundNode = false;
        for (NodeManager nodeManager : nodeManagers.values()) {
            log.debug("Checking node " + nodeManager.getNodeId() + " for resources");
            if (nodeManager.hasEnoughResource(requiredResource)) {
                foundNode = true;
                log.debug("Node " + nodeManager.getNodeId() + " has enough resources");
                // 创建容器
                String containerId = "container_" + nextContainerId++;
                Container container = new Container(
                    containerId,
                    nodeManager.getNodeId(),
                    requiredResource,
                    application.getApplicationId()
                );

                // 分配容器
                    log.debug("Attempting to allocate container " + containerId + " on node " + nodeManager.getNodeId());
                    boolean allocationResult = nodeManager.allocateContainer(container);
                    log.debug("Container allocation result: " + allocationResult);
                    
                    if (allocationResult) {
                        application.addContainer(container);
                        log.debug("Container " + containerId + " allocated to application " + application.getApplicationId());
                        application.setStatus(ApplicationMaster.ApplicationStatus.RUNNING);
                        log.debug("Application " + application.getApplicationId() + " status changed to RUNNING");

                        // 检查是否所有容器都已分配
                        allocatedContainers = application.getAllocatedContainers().size();
                        if (allocatedContainers >= numContainers) {
                            log.debug("All containers allocated for application " + application.getApplicationId() + ", removing from queue");
                            applicationQueue.poll();
                        }
                    } else {
                        log.debug("Failed to allocate container " + containerId + " on node " + nodeManager.getNodeId());
                        log.debug("Node status: running=" + nodeManager.isRunning());
                        log.debug("Node available resources: " + nodeManager.getAvailableResource().getMemoryMB() + "MB memory, " + nodeManager.getAvailableResource().getCpuCores() + " CPU cores");
                        log.debug("Node required resources: " + requiredResource.getMemoryMB() + "MB memory, " + requiredResource.getCpuCores() + " CPU cores");
                    }
                // 只分配一个容器，然后退出循环，让调度器在下一次运行时继续分配
                break;
            }
        }

        if (!foundNode) {
            log.debug("No node with enough resources found for application " + application.getApplicationId());
        }

        // 如果应用状态还是SUBMITTED，尝试再次设置为RUNNING（以防没有分配容器但应该运行的情况）
        if (application.getStatus() == ApplicationMaster.ApplicationStatus.SUBMITTED) {
            log.debug("Application " + application.getApplicationId() + " still SUBMITTED after scheduling attempt");
        }
    }

    /**
     * 完成应用程序
     */
    @Override
    public void completeApplication(String applicationId) {
        ApplicationMaster application = applications.get(applicationId);
        if (application != null) {
            log.debug("Completing application: " + applicationId + ", current status: " + application.getStatus());
            application.setStatus(ApplicationMaster.ApplicationStatus.COMPLETED);
            // 释放所有容器资源
            for (Container container : application.getAllocatedContainers()) {
                NodeManager nodeManager = nodeManagers.get(container.getNodeId());
                if (nodeManager != null) {
                    nodeManager.releaseContainer(container.getContainerId());
                    log.debug("Released container: " + container.getContainerId() + " for application: " + applicationId);
                }
            }
            log.debug("Application completed: " + applicationId);
        } else {
            log.debug("Application not found: " + applicationId);
        }
    }

    /**
     * 获取应用程序
     */
    @Override
    public ApplicationMaster getApplication(String applicationId) {
        return applications.get(applicationId);
    }

    @Override
    public Map<String, NodeManager> getNodeManagers() { return nodeManagers; }

    @Override
    public Map<String, ApplicationMaster> getApplications() { return applications; }
}